//
// Created by zhangcc on 2020/7/4.
//

#ifndef CRAFTNET_H
#define CRAFTNET_H

#include <NvInfer.h>
#include <NvOnnxParser.h>
#include <opencv2/opencv.hpp>

#include "utils/cudaWrapper.h"
#include "utils/ioHelper.h"

#define MAX_WORKSPACE_SIZE ((1ULL << 30) * 13) // 1 GB
#define MAX_BINDING_NUMS   (2)

class AgnleNet {
public:
    AgnleNet(int device_id, bool trt_serialize) :
            device_id(device_id), trt_serialize(trt_serialize) {
    };

    int init(const std::string &model_file);

    std::vector<float> extract(const cv::Mat &inputs);

    ~AgnleNet();

private:

    int serialization(const std::string &model_file, const std::string &engine_file);

    bool Input(const cv::Mat &src, float *data);

private:
    std::unique_ptr <nvinfer1::ICudaEngine, nvinfer1::Destroy<nvinfer1::ICudaEngine>> mEngine{nullptr};
    std::unique_ptr <nvinfer1::IExecutionContext, nvinfer1::Destroy<nvinfer1::IExecutionContext>> context{nullptr};

    void *bindings[MAX_BINDING_NUMS]{0};
    cudaStream_t stream_t;

    int inputBufferSize;
    int outputBufferSize;

    int device_id = 0;
    int max_batch_size = 1;
    int channels = 3;
    bool trt_serialize;
};

#endif //CRAFTNET_H
